# plot_result.py (for PVPKEET)
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

def plot_pvpkeet_results(df):
    """
    Reads CSV data for the PVPKEET scheme and generates enhanced performance plots,
    showing both average and total times for Test, Proof, and Verify.
    """
    sns.set_theme(style="whitegrid")

    # --- 数据预处理 ---
    df.columns = df.columns.str.strip()
    
    df_filesize = df[df['Experiment'] == 'FileSize'].copy().sort_values(by='Value').reset_index(drop=True)
    df_cspcount = df[df['Experiment'] == 'CSPCount'].copy().sort_values(by='Value').reset_index(drop=True)
    
    # --- 图表绘制 ---
    fig, axes = plt.subplots(2, 2, figsize=(20, 16))
    fig.suptitle('Performance Analysis of PVPKEET Scheme', fontsize=22, fontweight='bold', y=0.97)
    
    # --- 实验 1: 文件大小 ---
    ax1_avg = axes[0, 0]
    ax1_total = axes[0, 1]
    
    # 平均时间图
    df_filesize.plot(x='Value', y=['SetupTime(ms)', 'EncryptTime(ms)', 'AvgTestTime(ms)', 'AvgProofTime(ms)', 'AvgVerifyTime(ms)'], 
                     marker='o', ax=ax1_avg, ms=8, lw=2.5)
    ax1_avg.set_title('Avg. Time vs. File Size', fontsize=16, fontweight='bold')
    ax1_avg.set_xlabel('File Size (MB)', fontsize=14)
    ax1_avg.set_ylabel('Average Time (ms)', fontsize=14)
    ax1_avg.legend(['Setup (Total)', 'Encrypt (Avg)', 'Test (Avg Pair)', 'Proof (Avg Pair)', 'Verify (Avg Pair)'], fontsize=12)
    
    # 总时间图
    df_filesize.plot(x='Value', y=['SetupTime(ms)', 'EncryptTime(ms)', 'TotalTestTime(ms)', 'TotalProofTime(ms)', 'TotalVerifyTime(ms)'], 
                     marker='o', ax=ax1_total, ms=8, lw=2.5)
    ax1_total.set_title('Total Time vs. File Size', fontsize=16, fontweight='bold')
    ax1_total.set_xlabel('File Size (MB)', fontsize=14)
    ax1_total.set_ylabel('Total Time (ms)', fontsize=14)
    ax1_total.legend(['Setup (Total)', 'Encrypt (Avg)', 'Test (Total)', 'Proof (Total)', 'Verify (Total)'], fontsize=12)

    # --- 实验 2: CSP 数量 ---
    ax2_avg = axes[1, 0]
    ax2_total = axes[1, 1]

    # 平均时间图
    df_cspcount.plot(x='Value', y=['SetupTime(ms)', 'EncryptTime(ms)', 'AvgTestTime(ms)', 'AvgProofTime(ms)', 'AvgVerifyTime(ms)'],
                     marker='s', ax=ax2_avg, ms=8, lw=2.5)
    ax2_avg.set_title('Avg. Time vs. Number of CSPs', fontsize=16, fontweight='bold')
    ax2_avg.set_xlabel('Number of CSPs (K)', fontsize=14)
    ax2_avg.set_ylabel('Average Time (ms)', fontsize=14)
    ax2_avg.legend(['Setup (Total)', 'Encrypt (Avg)', 'Test (Avg Pair)', 'Proof (Avg Pair)', 'Verify (Avg Pair)'], fontsize=12)

    # 总时间图
    df_cspcount.plot(x='Value', y=['SetupTime(ms)', 'EncryptTime(ms)', 'TotalTestTime(ms)', 'TotalProofTime(ms)', 'TotalVerifyTime(ms)'],
                     marker='s', ax=ax2_total, ms=8, lw=2.5)
    ax2_total.set_title('Total Time vs. Number of CSPs', fontsize=16, fontweight='bold')
    ax2_total.set_xlabel('Number of CSPs (K)', fontsize=14)
    ax2_total.set_ylabel('Total Time (ms)', fontsize=14)
    ax2_total.legend(['Setup (Total)', 'Encrypt (Avg)', 'Test (Total)', 'Proof (Total)', 'Verify (Total)'], fontsize=12)
    
    # --- 统一调整 ---
    for ax in axes.flatten():
        ax.tick_params(axis='both', which='major', labelsize=12)
        y_max = ax.get_ylim()[1]
        if y_max > 5000:
             ax.set_yscale('log')
             ax.set_ylabel(ax.get_ylabel() + ' (log scale)', fontsize=14)

    plt.tight_layout(rect=[0, 0, 1, 0.96])
    plt.savefig('performance_plots_pvpkeet.png', dpi=300, bbox_inches='tight')
    print("Performance plots for PVPKEET saved to 'performance_plots_pvpkeet.png'")
    plt.show()

if __name__ == '__main__':
    csv_filename = 'results_pvpkeet.csv'
    try:
        df = pd.read_csv(csv_filename)
        plot_pvpkeet_results(df)
    except FileNotFoundError:
        print(f"Error: '{csv_filename}' not found. Please run the Go program first.")
    except Exception as e:
        print(f"An error occurred: {e}")